import numpy as np
import random
import matplotlib.pyplot as plt


class BSO(object):

    def __init__(self, obj, n, dim=2, search_domain=[], nc=5, iters=20, P_ucc=0.9, P_gi=0.1, P_sc=0.5, P_scc=0.5, P_2scc=0.5):
        self.obj = obj  # objection function
        self.n = n  # number potential solution
        self.dim = dim  # dimension of point
        self.nc = nc  # number of clusters
        self.iters = iters  # number iteration
        self.P_ucc = P_ucc  # update cluster center probability
        self.P_gi = P_gi  # generate individual probability
        self.P_sc = P_sc  # selected cluster probability
        self.P_scc = P_scc  # selected cluster center probability
        self.P_2scc = P_2scc  # two selected cluster center probability
        # search domain
        if len(search_domain) == 0: self.search_domain = [[0, 1] for _ in range(dim)]
        else: self.search_domain = search_domain

    def init_sols(self, n):
        # init n sols
        sols = []
        for i in range(n):
            # the uniform distribution domain
            sols.append([random.uniform(self.search_domain[_][0], self.search_domain[_][1]) for _ in range(self.dim)])
        return np.array(sols)

    def cluster(self, centers, sets):
        # use kmeans method to cluster sols
        from kmeans import kmeans
        return kmeans(centers, sets)

    def process(self):
        fig = plt.figure(figsize=(10, 10))
        plt.ion()
        X = np.arange(-2.5, 2.5, 0.05)
        Y = np.arange(-2.5, 2.5, 0.05)
        X, Y = np.meshgrid(X, Y)
        Z = self.obj(np.array([X, Y]))
        cs = plt.contourf(X, Y, Z)
        fig.colorbar(cs)

        n_sols = self.init_sols(self.n)

        for i in range(self.iters):
            # print(n_sols)
            s = plt.scatter(n_sols[:, 0], n_sols[:, 1], color='blue')

            # cluster n sols into m clusters
            m_center = self.init_sols(self.nc)

            colors = ['red', 'blue', 'green', 'orange', 'gray']
            colors = colors[0:len(m_center)]
            z = plt.scatter(m_center[:, 0], m_center[:, 1], color=colors)
            plt.pause(0.5)
            s.remove()
            z.remove()

            m_center, clusters = self.cluster(m_center, n_sols)
            colors = colors[0:len(m_center)]
            z = plt.scatter(m_center[:, 0], m_center[:, 1], s=100, color=colors, marker='*')
            ic = 0
            s = []
            for c in clusters:
                s.append(plt.scatter(c[:, 0], c[:, 1], s=100, color=colors[ic]))
                ic += 1
            plt.pause(0.8)
            for item in s:
                item.remove()
            z.remove()

            new_n_sols = []
            compare_idx = 0

            # evaluate the n sols
            new_centers = []
            for items in clusters:
                ranks = []
                for _ in range(len(items)):
                    ranks.append(self.obj(items[_]))
                # get the fitness of obj at cluster
                idx = np.argmin(np.array(ranks))
                ranks.clear()
                new_centers.append(items[idx])
                new_n_sols.append(items[idx])
                compare_idx += 1

            r_1 = random.random()
            if r_1 < self.P_ucc:
                idx = random.randint(0, len(new_centers)-1)
                new_sol = self.init_sols(1)[0]
                new_centers[idx] = list(new_sol)

            # generates n new sols
            while len(new_n_sols) < self.n:
                r_2 = random.random()
                # new_sol = None
                if r_2 < self.P_gi:
                    idx = random.randint(0, len(new_centers)-1)
                    r_3 = random.random()
                    if r_3 < self.P_scc:
                        new_sol = [ax+random.uniform(-1.5, 1.5) for ax in new_centers[idx]]
                        # limit the range of sol
                        for k in range(len(new_sol)):
                            if new_sol[k] > self.search_domain[k][1]:
                                new_sol[k] = self.search_domain[k][1]
                            elif new_sol[k] < self.search_domain[k][0]:
                                new_sol[k] = self.search_domain[k][0]
                    else:
                        idx_c = random.randint(0, len(clusters[idx])-1)
                        new_sol = [ax + random.uniform(-1.5, 1.5) for ax in clusters[idx][idx_c]]
                        # limit the range of sol
                        for k in range(len(new_sol)):
                            if new_sol[k] > self.search_domain[k][1]:
                                new_sol[k] = self.search_domain[k][1]
                            elif new_sol[k] < self.search_domain[k][0]:
                                new_sol[k] = self.search_domain[k][0]
                else:
                    idx1, idx2 = random.sample(range(len(new_centers)), 2)
                    r_4 = random.random()
                    if r_4 < self.P_2scc:
                        new_sol = [(new_centers[idx1][k] + new_centers[idx2][k])/2 + random.uniform(-1.5, 1.5)
                                   for k in range(len(new_centers[idx1]))]
                        # limit the range of sol
                        for k in range(len(new_sol)):
                            if new_sol[k] > self.search_domain[k][1]:
                                new_sol[k] = self.search_domain[k][1]
                            elif new_sol[k] < self.search_domain[k][0]:
                                new_sol[k] = self.search_domain[k][0]
                    else:
                        idx1_c = random.randint(0, len(clusters[idx1])-1)
                        idx2_c = random.randint(0, len(clusters[idx2])-1)
                        new_sol = [(clusters[idx1][idx1_c][k] + clusters[idx2][idx2_c][k])/2 + random.uniform(-1.5, 1.5)
                                   for k in range(len(new_centers[idx1]))]
                        # limit the range of sol
                        for k in range(len(new_sol)):
                            if new_sol[k] > self.search_domain[k][1]:
                                new_sol[k] = self.search_domain[k][1]
                            elif new_sol[k] < self.search_domain[k][0]:
                                new_sol[k] = self.search_domain[k][0]
                if self.obj(np.array(new_sol)) > self.obj(n_sols[compare_idx]):
                    new_n_sols.append(np.array(new_sol))
                else:
                    new_n_sols.append(n_sols[compare_idx])
                compare_idx += 1
            n_sols = np.array(new_n_sols)

        best = n_sols[0]
        for i in range(len(n_sols)):
            if self.obj(best) > self.obj(n_sols[i]):
                best = n_sols[i]
        fitness = self.obj(best)
        plt.scatter(best[0], best[1], color='red', marker='v')
        strs = '({}, {}, {})'.format(round(best[0], 2), round(best[1], 2), round(fitness, 2))
        plt.annotate(strs, xy=(best[0], best[1]), xytext=(best[0]+0.1, best[1]+0.1))
        plt.ioff()
        plt.show()


def f(x):
    return np.sin(np.sqrt(np.add.reduce(np.power(x, 2))))/np.sqrt(np.add.reduce(np.power(x, 2))) + \
           np.exp(np.add.reduce(np.cos(2*np.pi*x))/2) - 2.71289


if __name__ == '__main__':
    bso = BSO(f, 50, 2, [[-2., 2.], [-2., 2.]])
    print(bso.process())
